import struct
import logging

import asyncio
from asyncio import Queue, StreamReader

from message import *

class ProtocolError(BaseException):
    pass

class PeerConnection:
    def __init__(self, queue: Queue, info_hash, peer_id, piece_manager, on_block_cb) -> None:
        self.my_state = []
        self.peer_state = []
        self.queue = queue
        self.info_hash = info_hash
        self.peer_id = peer_id
        self.remote_id = None
        self.writer = None
        self.reader = None
        self.piece_manager = piece_manager
        self.on_block_cb = on_block_cb
        self.future = asyncio.ensure_future(self._start())
        
    async def _start(self):
        while 'stopped' not in self.my_state:
            ip, port = await self.queue.get()
            logging.info('Deal peer: {}:{}'.format(ip, port))

            try:
                fut = asyncio.open_connection(ip, port)
                self.reader, self.writer = await asyncio.wait_for(fut, timeout=10)
                # self.reader, self.writer = await asyncio.open_connection(ip, port)
                logging.info('Connection to peer: {}:{}'.format(ip, port))

                buffer = await self._send_handshake()
                # TODO: need verify handshake 

                # The default state for a connection is that peer is not
                # interested and we are choked
                self.my_state.append('choked')

                await self._send_interested()
                self.my_state.append('interested')

                async for message in PeerStreamIterator(self.reader, buffer):
                    if 'stopped' in self.my_state:
                        break
                    
                    logging.debug('Peer said: {}'.format(message))
                    if type(message) is BitField:
                        self.piece_manager.add_peer(self.remote_id, message.bitfield)
                    elif type(message) is Interested:
                        self.peer_state.append('interested')
                    elif type(message) is NotInterested:
                        self.peer_state.remove('interested')
                    elif type(message) is Choke:
                        self.my_state.append('choked')
                    elif type(message) is Unchoke:
                        if 'choked' in self.my_state:
                            self.my_state.remove('choked')
                    elif type(message) is Have:
                        self.piece_manager.update_peer(self.remote_id, message.index)
                    elif type(message) is KeepAlive:
                        pass
                    elif type(message) is Piece:
                        self.my_state.remove('pending_request')
                        self.on_block_cb(
                            peer_id=self.remote_id,
                            piece_index=message.index,
                            block_offset=message.begin,
                            data=message.block
                        )
                    elif type(message) is Request:
                        # TODO: Request msg
                        logging.info('TODO: Request msg')
                    elif type(message) is Cancel:
                        logging.info('TODO: Cancel msg')
                    
                    if 'choked' not in self.my_state:
                        if 'interested' in self.my_state:
                            self.my_state.append('pending_request')
                            await self._request_piece()
                    
                    
            except ProtocolError as e:
                logging.info(e)
            except asyncio.TimeoutError as e:
                logging.info('Timeout peer: {}:{}'.format(ip, port))
            except TimeoutError as e:
                logging.info('Timeout peer: {}:{}'.format(ip, port))
            except ConnectionRefusedError as e:
                logging.info('Connection refused peer: {}:{}'.format(ip, port))
            except OSError as e:
                logging.info('Connect failed peer: {}:{}'.format(ip, port))
                # logging.error(e)
            except BaseException as e:
                # logging.error(type(e))
                logging.exception(e)
    
    async def _send_handshake(self):
        logging.debug('Send handshake')
        self.writer.write(Handshake(self.info_hash, self.peer_id).encode())
        await self.writer.drain()

        buf = b''
        tries = 0
        while len(buf) < Handshake.HANDSHAKE_LENGTH and tries < 10:
            tries += 1
            buf = await self.reader.read(PeerStreamIterator.CHUNK_SIZE)
        
        if len(buf) == 0:
            raise ProtocolError('Recv 0 size handshake')
        
        logging.debug('Got responsed({}): {}'.format(len(buf), buf))
        # must slice, the handshake recv may have bitfield
        response = Handshake.decode(buf[:Handshake.HANDSHAKE_LENGTH])
        logging.info('Got responsed: {}'.format(response))

        # TODO: According to spec we should validate that the peer_id received
        # from the peer match the peer_id received from the tracker.
        self.remote_id = response.peer_id
        logging.info('Handshake with peer {} was successful'.format(response.peer_id))

        # We need to return the remaining buffer data, since we might have
        # read more bytes then the size of the handshake message and we need
        # those bytes to parse the next message.
        return buf[Handshake.HANDSHAKE_LENGTH:]

    async def _send_interested(self):
        message = Interested().encode()
        logging.debug('Send message: Interested')
        self.writer.write(message)
        await self.writer.drain()
    
    def stop(self):
        self.my_state.append('stopped')
        if not self.future.done():
            self.future.cancel()
    
    async def _request_piece(self):
        block = self.piece_manager.next_request(self.remote_id)
        if block:
            message = Request(block.piece, block.offset, block.length).encode()

            logging.debug('Requesting block({}) {} of piece {} from peer {}'.format(
                block.length,
                block.offset,
                block.piece,
                self.remote_id
            ))
            
            self.writer.write(message)
            await self.writer.drain()

class PeerStreamIterator:

    CHUNK_SIZE = 10*1024

    def __init__(self, reader: StreamReader, initial: bytes=None) -> None:
        self.reader = reader
        self.buffer = initial if initial else b''

    def __aiter__(self):
        return self

    async def __anext__(self):
        while True:
            try:
                if self.buffer:
                    message = self.parse()
                    if message:
                        return message
                    
                data = await self.reader.read(PeerStreamIterator.CHUNK_SIZE)
                if data:
                    self.buffer += data
                else:
                    logging.debug('Read FIN of connection')
                    if self.buffer:
                        message = self.parse()
                        if message:
                            return message
                    raise StopAsyncIteration()
                
            except ConnectionResetError:
                logging.debug('Connection closed by peer')
                raise StopAsyncIteration()
            except StopAsyncIteration:
                raise StopAsyncIteration()
            except Exception as e:
                logging.exception(e)
                raise StopAsyncIteration()
    
    def parse(self):
        """
        Tries to parse protocol messages if there is enough bytes read in the
        buffer.

        :return The parsed message, or None if no message could be parsed
        """
        # Each message is structured as:
        #     <length prefix><message ID><payload>
        #
        # The `length prefix` is a four byte big-endian value
        # The `message ID` is a decimal byte
        # The `payload` is the value of `length prefix`
        #
        # The message length is not part of the actual length. So another
        # 4 bytes needs to be included when slicing the buffer.
        header_length = 4

        if len(self.buffer) > 4:
            message_length = struct.unpack('>I', self.buffer[:4])[0]

            if message_length == 0:
                pass # keep alive

            if len(self.buffer) >= message_length:
                _, message_id = struct.unpack('>Ib', self.buffer[:5])

                def _consume():
                    # eat current message
                    self.buffer = self.buffer[header_length + message_length:]
                
                def _data():
                    return self.buffer[:header_length + message_length]
                
                if message_id is PeerMessage.KeepAlive:
                    data = _data()
                    _consume()
                    return KeepAlive.decode(data)
                elif message_id is PeerMessage.Choke:
                    data = _data()
                    _consume()
                    return Choke.decode(data)
                elif message_id is PeerMessage.Unchoke:
                    data = _data()
                    _consume()
                    return Unchoke.decode(data)
                elif message_id is PeerMessage.Interested:
                    data = _data()
                    _consume()
                    return Interested.decode(data)
                elif message_id is PeerMessage.NotInterested:
                    data = _data()
                    _consume()
                    return NotInterested.decode(data)
                elif message_id is PeerMessage.Have:
                    data = _data()
                    _consume()
                    return Have.decode(data)
                elif message_id is PeerMessage.BitField:
                    data = _data()
                    _consume()
                    return BitField.decode(data)
                elif message_id is PeerMessage.Request:
                    data = _data()
                    _consume()
                    return Request.decode(data)
                elif message_id is PeerMessage.Piece:
                    data = _data()
                    _consume()
                    return Piece.decode(data)
                elif message_id is PeerMessage.Cancel:
                    data = _data()
                    _consume()
                    return Cancel.decode(data)
                
                logging.error('Message({}) with id {} NOT match all\n{}'.format(len(self.buffer), message_id, self.buffer.hex()))
                
                
                
                
                